import base64
import io
import os
import sys

import gradio as gr
from PIL import ImageDraw, ImageFont

from sdk.httpclient import CTClientBuilder, CTClient

base_url = "https://ai-global.ctapi.ctyun.cn"

client: CTClient


def _get_image_base64(img):
    """
    返回图片的base64大小（单位MB）和内容
    """
    byte_arr = io.BytesIO()
    img.save(byte_arr, format='PNG')
    byte_arr = byte_arr.getvalue()
    base64_str = base64.b64encode(byte_arr).decode('utf-8')
    base64_mb_size = sys.getsizeof(base64_str) / (1024 * 1024)
    return base64_mb_size, base64_str


def draw_faces(image, face_locations):
    draw = ImageDraw.Draw(image)

    # 字体
    font_path = './font/NotoSans-Regular.ttf'  # 请确保这个路径正确
    font_info = ImageFont.truetype(font_path, size=30)  # 设置字体大小

    idx=1
    for face in face_locations:
        # 获取人脸位置信息
        left = face['face_location']['left']
        top = face['face_location']['top']
        width = face['face_location']['width']
        height = face['face_location']['height']
        
        # 戴口罩场景，仅显示满足 戴口罩 条件
        if face.get('Mask') is not None:
            if not face.get('Mask'):
                continue
        # 人脸活体检测场景，仅显示满足 活体人脸 条件
        elif face.get('FaceAntiSpoofing') is not None:
            if face.get('FaceAntiSpoofing') != '活体人脸':
                continue
        
        # 其它场景，默认显示
        # 绘制矩形框
        draw.rectangle([left, top, left + width, top + height], outline="red", width=2)

        # 绘制文本
        text = str(idx)
        # 计算文本的位置
        text_bbox = draw.textbbox((0, 0), text, font=font_info)
        text_width = text_bbox[2] - text_bbox[0]
        text_height = text_bbox[3] - text_bbox[1]
        # 计算文本的位置，使得文本紧挨着矩形框的下方
        text_left = left + (width - text_width) // 2  # 文本居中对齐
        text_top = top + height + 10  # 文本距离矩形框下方10像素
        # 绘制文本，使用红色字体
        draw.text((text_left, text_top), text, fill="red", font=font_info)
        idx = idx + 1
    return image


def detect_face(img):
    base64_size, base64_str = _get_image_base64(img)
    if base64_size > 2.0:
        return {"errMsg": f"上传的图片进行base64编码后大小为 {base64_size:.2f} MB，超过 2.00 MB"}, None

    url = f"{base_url}/v1/aiop/api/2f6hqix09mv4/face/PERSON/person/detectFaceFromBase64"
    response = client.request("POST", url, body={"imageContent": base64_str})
    response.raise_for_status()
    response_json = response.json()
    try:
        result = response_json['returnObj']
        return draw_faces(img, result["face_list"]), result
    except Exception as e:
        print(f"error: {e}")
        return None, response_json


def detect_age_gender(img):
    base64_size, base64_str = _get_image_base64(img)
    if base64_size > 2.0:
        return {"errMsg": f"上传的图片进行base64编码后大小为 {base64_size:.2f} MB，超过 2.00 MB"}, None

    url = f"{base_url}/v1/aiop/api/2f6hw5o5t7gg/face/PERSON/person/detectAgeGenderFromBase64"
    response = client.request("POST", url, body={"imageContent": base64_str})
    response.raise_for_status()
    response_json = response.json()
    try:
        result = response_json['returnObj']
        return draw_faces(img, result["face_list"]), result
    except Exception as e:
        print(f"error: {e}")
        return None, response_json


def compare_face(img1, img2):
    base64_size1, base64_str1 = _get_image_base64(img1)
    base64_size2, base64_str2 = _get_image_base64(img2)
    print(f"{base64_size1}, {base64_size2}")
    if base64_size1 > 2.0 or base64_size2 > 2.0:
        return {
            "errMsg": f"上传的图片进行base64编码后大小为 {base64_size1:.2f} MB、{base64_size2:.2f} MB，超过 2.00 MB"}, None

    url = f"{base_url}/v1/aiop/api/2f7awxekgvls/face/compare/PERSON/person/compareFromBase64"
    response = client.request("POST", url, body={"img1Base64": base64_str1, "img2Base64": base64_str2})
    response.raise_for_status()
    response_json = response.json()
    try:
        result = response_json['returnObj']
        return result
    except Exception as e:
        return response_json


def detect_fas(img):
    base64_size, base64_str = _get_image_base64(img)
    if base64_size > 2.0:
        return {"errMsg": f"上传的图片进行base64编码后大小为 {base64_size:.2f} MB，超过 2.00 MB"}, None

    url = f"{base_url}/v1/aiop/api/2hfksnibjaos/face-fas-action/person/detectFasFromBase64"
    response = client.request("POST", url, body={"imageContent": base64_str})
    response.raise_for_status()
    response_json = response.json()
    try:
        result = response_json['returnObj']
        return draw_faces(img, result["face_list"]), result
    except Exception as e:
        print(f"error: {e}")
        return None, response_json


def detect_mask(img):
    base64_size, base64_str = _get_image_base64(img)
    if base64_size > 2.0:
        return {"errMsg": f"上传的图片进行base64编码后大小为 {base64_size:.2f} MB，超过 2.00 MB"}, None

    url = f"{base_url}/v1/aiop/api/2f6hycj3a9z4/face/PERSON/person/detectMaskFromBase64"
    response = client.request("POST", url, body={"imageContent": base64_str})
    response.raise_for_status()
    response_json = response.json()
    try:
        result = response_json['returnObj']
        return draw_faces(img, result["face_list"]), result
    except Exception as e:
        print(f"error: {e}")
        return None, response_json


product_intro = "进一步了解天翼云人脸识别产品：<a href='https://www.ctyun.cn/products/facerecognition'>https://www.ctyun.cn/products/facerecognition</a>"
gr_config = {
    "detect_face": gr.Interface(fn=detect_face,
                                inputs=gr.Image(type="pil", label="上传图像"),
                                outputs=[
                                    gr.Image(type="pil", label="预览"),
                                    gr.Textbox(label="检测结果")
                                ],
                                examples=[],
                                title="人脸识别 - 人脸检测",
                                description="用于检测输入图像中的人脸，输出人脸位置坐标。",
                                article=product_intro),
    "detect_age_gender": gr.Interface(fn=detect_age_gender,
                                      inputs=gr.Image(type="pil", label="上传图像"),
                                      outputs=[
                                          gr.Image(type="pil", label="预览"),
                                          gr.Textbox(label="识别结果")
                                      ],
                                      examples=[],
                                      title="人脸识别 - 人脸属性识别",
                                      description="用于检测输入图像中的人脸年龄、性别等属性。",
                                      article=product_intro),
    "compare_face": gr.Interface(fn=compare_face,
                                 inputs=[
                                     gr.Image(type="pil", label="上传图像1"),
                                     gr.Image(type="pil", label="上传图像2")
                                 ],
                                 outputs=[
                                     gr.Textbox(label="比对结果")
                                 ],
                                 examples=[],
                                 title="人脸识别 - 人脸比对",
                                 description="用于检测输入的两张图像中的人脸相似度。",
                                 article=product_intro),
    "detect_fas": gr.Interface(fn=detect_fas,
                               inputs=[
                                   gr.Image(type="pil", label="上传图像")
                               ],
                               outputs=[
                                   gr.Image(type="pil", label="预览"),
                                   gr.Textbox(label="检测结果")
                               ],
                               examples=[],
                               title="人脸识别 - 人脸活体检测",
                               description="用于检测输入图像中的人脸是否为活体。",
                               article=product_intro),
    "detect_mask": gr.Interface(fn=detect_mask,
                                inputs=[
                                    gr.Image(type="pil", label="上传图像")
                                ],
                                outputs=[
                                    gr.Image(type="pil", label="预览"),
                                    gr.Textbox(label="识别结果")
                                ],
                                examples=[],
                                title="人脸识别 - 是否戴口罩识别",
                                description="用于检测输入图像中的人脸是否戴口罩。",
                                article=product_intro),
}


def get_not_empty_env(key):
    value = os.getenv(key, "").strip()
    if value == "":
        raise Exception(f"env {key} is not set or empty")
    return value


def run():
    ctyun_ak = get_not_empty_env("ext_cf_ctyun_ak")
    ctyun_sk = get_not_empty_env("ext_cf_ctyun_sk")
    ctyun_ai_app_key = get_not_empty_env("ext_cf_ctyun_ai_app_key")
    fr_type = get_not_empty_env("ext_cf_fr_type")

    global client
    client = CTClientBuilder().with_ak(ctyun_ak).with_sk(ctyun_sk).with_ai_app_key(ctyun_ai_app_key).build()

    gr_config[fr_type].launch(server_name="0.0.0.0", server_port=9000)


if __name__ == '__main__':
    run()
